import random
import time

import numpy as np
import cv2

from cjc import project_setting


def calculate_position(xywh):
    """根据底部坐标判断是否出现过"""
    return [int((xywh[0] + xywh[2]) / 2), int(xywh[3])]


class Trajectories:
    '''
    用来画人的轨迹图,
    第一步创建画布  self.reset_canvas(width, height)  或者 create_canvas_from_bk
    第二部  画点  draw_point_with_index  或者draw_point
    '''

    def __init__(self, color_type=0, width=None, height=None):
        '''
        color_type 0 是float32类型，1 是unit8类型
        Parameters
        ----------
        color_type
        width
        height
        '''
        self.color_type = color_type
        self.traj_map = None
        self.point_size = 1
        self.point_color = (0, 0, 1)  # (1, 1, 1)
        self.thickness = 4
        if width is not None and height is not None:
            self.reset_canvas(width, height)
        self.background = None
        self.point_colors = [self.rand_color() for _ in range(20)]
        self.start_time = None  #
        self.indexed_time = {}

    def clear(self):
        self.reset_canvas()

    def get_point_color(self, index):
        if self.point_colors is None or len(self.point_colors) == 0:
            self.refresh_color()
            return self.point_color
        mod_index = int(index % len(self.point_colors))
        return self.point_colors[mod_index]

    def create_canvas_from_bk(self, im0):
        """实际使用的时候是用这个函数初始化画布"""
        if self.background is None:
            self.set_background(im0)
            w, h = im0.shape[1], im0.shape[0]
            self.reset_canvas(w, h)

    def set_background(self, bk):
        """设置合成背景图"""

        if bk[0][0][0] > 1:
            bk = bk.astype(np.float32) / 255
        self.background = bk
        # print(bk[0][0][0])
        # print('更换了背景')

    @staticmethod
    def to_float32(img):
        if len(img.shape) >= 3:
            if img[0][0][0] > 1:
                return img.astype(np.float32) / 255
        return img
    @staticmethod
    def to_uint8(img):
        if len(img.shape) >= 3:
            if img[0][0][0] <=1:
                return (img* 255).astype("uint8")
        return img
    def reset_canvas(self, width=None, height=None):
        c = 3
        if width is None or height is None:
            height, width, c = self.traj_map.shape
        # print(width,height,c)
        if self.color_type == 0:
            self.traj_map = np.zeros((height, width, c)).astype(
                np.float32)  # 参照https://blog.csdn.net/qq_40985985/article/details/116025825
        else:
            self.traj_map = np.zeros((height, width, c)).astype(
                'uint8')  # 参照https://blog.csdn.net/qq_40985985/article/details/116025825

    def draw_foot_point(self, xywh):
        """根据底部坐标判断是否出现过"""
        p = [int((xywh[0] + xywh[2]) / 2), int(xywh[3])]
        self.draw_point(p)

    def draw_point(self, cord):
        """画点"""
        cv2.circle(self.traj_map, cord, self.point_size, self.point_color,
                   self.thickness)  # 参照https://blog.csdn.net/zxs0222/article/details/107427556

    def set_last_draw_time(self, index):
        self.indexed_time[index] = time.time()

    def can_draw_indexed_point(self,index):
        if index in self.indexed_time:
            return time.time() - self.indexed_time[index] > project_setting.traj_map_interval
        return True

    def draw_point_with_index(self, cord, index):
        """画点"""
        if self.can_draw_indexed_point(index):
            c = self.get_point_color(index)
            xy = calculate_position(cord)
            cv2.circle(self.traj_map, xy, self.point_size, c,
                       self.thickness)  # 参照https://blog.csdn.net/zxs0222/article/details/107427556
            self.set_last_draw_time(index)

    def show_point(self):
        cv2.imshow('Trajectories', self.traj_map)
        cv2.waitKey(5000)

    def show_result(self):
        result = self.result_with_background()
        cv2.imshow('Trajectories', result)
        cv2.waitKey(5000)

    def rand01(self):
        if self.color_type == 0:
            return random.random()
        else:
            return random.randrange(0, 255)

    def rand_color(self):
        return self.rand01(), self.rand01(), self.rand01()

    def refresh_color(self):
        self.point_color = (self.rand01(), self.rand01(), self.rand01())

    def result(self):
        return self.traj_map

    def uint8_result(self):
        return self.to_uint8(self.traj_map)

    def result_to_tracker(self):
        mp = self.result_with_background()
        return (mp * 255).astype('uint8')

    def combine_traj_and_img(self, img):
        img = self.to_float32(img)
        img2gray = cv2.cvtColor(self.traj_map, cv2.COLOR_BGR2GRAY)
        # cv2.threshold (源图片, 阈值=ret, 填充色, 阈值类型)
        ret, mask = cv2.threshold(img2gray, 0.01, 1.0, cv2.THRESH_BINARY)
        mask = (mask * 255).astype("uint8")
        mask_inv = cv2.bitwise_not(mask)
        mask_inv = mask_inv.astype("uint8")  # 参照  https://cloud.tencent.com/developer/ask/sof/594717
        img1_bg = cv2.bitwise_and(img, img, mask=mask_inv)
        img1_bg = img1_bg.astype('float32')
        dst = cv2.add(img1_bg, self.traj_map)
        return (dst * 255).astype('uint8')

    def result_with_background(self):
        """参照https://blog.csdn.net/MasterCayman/article/details/121389609
           参照2  https://blog.csdn.net/weixin_45335726/article/details/122415833"""
        # img1 = cv2.resize(img1, (600, 300))  # 变换尺寸，这里用不到。
        img2gray = cv2.cvtColor(self.traj_map, cv2.COLOR_BGR2GRAY)
        # cv2.threshold (源图片, 阈值=ret, 填充色, 阈值类型)
        ret, mask = cv2.threshold(img2gray, 0.01, 1.0, cv2.THRESH_BINARY)
        mask = (mask * 255).astype("uint8")
        mask_inv = cv2.bitwise_not(mask)
        mask_inv = mask_inv.astype("uint8")  # 参照  https://cloud.tencent.com/developer/ask/sof/594717
        img1_bg = cv2.bitwise_and(self.background, self.background, mask=mask_inv)
        img1_bg = img1_bg.astype('float32')
        dst = cv2.add(img1_bg, self.traj_map)
        return dst


if __name__ == '__main__':
    f = r'C:\Users\cjc.win\Pictures\8594-63f26980881459683dc334a45e52a472.jpg'
    # mp = cv2.imread(f).astype(np.float)
    mp = cv2.imread(f)
    t = Trajectories(0)
    # t.set_background(mp)
    t.create_canvas_from_bk(mp)
    t.refresh_color()
    t.draw_point_with_index([222, 222, 56, 56], 21)
    t.draw_point_with_index([338, 12, 56, 56], 21)
    t.draw_point_with_index([232, 12, 89, 90], 9)
    t.clear()
    t.draw_point_with_index([158, 12, 256, 56], 21)
    t.draw_point((20, 20))
    t.draw_point((76, 23))
    # t.draw_point((122, 34))
    # t.draw_point((160, 480))
    t.show_result()
    cv2.waitKey(3000)
